"""
基础研究工作流
定义研究助手的基本工作流程
"""

from langgraph.graph import StateGraph, END
from typing import AsyncGenerator, Dict, Any
from src.research_core.state import AgentState

def create_workflow():
    """创建工作流"""
    from src.research_core.nodes import (
        generate_search_query,
        execute_web_search,
        evaluate_results,
        generate_final_answer,
        should_continue
    )
    
    # 创建状态图
    workflow = StateGraph(AgentState)

    # 定义适配器函数来处理类型转换
    def generate_query_node(state):
        return generate_search_query(state)
    
    def web_search_node(state):
        return execute_web_search(state)
    
    def evaluate_node(state):
        return evaluate_results(state)
    
    def generate_answer_node(state):
        return generate_final_answer(state)

    # 添加节点
    workflow.add_node("generate_query", generate_query_node)   # 添加生成查询节点
    workflow.add_node("web_search", web_search_node)          # 添加网络搜索节点
    workflow.add_node("evaluate", evaluate_node)              # 添加评估节点
    workflow.add_node("generate_answer", generate_answer_node)  # 添加生成答案节点

    # 设置入口点
    workflow.set_entry_point("generate_query")

    # 添加边
    workflow.add_edge("generate_query", "web_search")
    workflow.add_edge("web_search", "evaluate")

    # 添加条件边
    workflow.add_conditional_edges(
        "evaluate",
        should_continue,
        {
            "end": "generate_answer",  # 如果完成，去生成答案
            "continue": "generate_query"  # 如果继续，回到生成查询
        }
    )

    # 添加最后的边
    workflow.add_edge("generate_answer", END)

    return workflow